import torch
from torch.utils.tensorboard import SummaryWriter

from datasets import FoodDataset

DatasetDir = "/mnt/DISK/xjk/pytorch_learning/datasets"

class Config:
    seeds=114514
    batch_size = 64
    epochs = 10000
    learning_rate = 0.001
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    checkpoint_dict_keys = ["model_dict", "optimizer_dict", "scheduler_dict"]
    checkpoint_dir = "/mnt/DISK/xjk/pytorch_learning/checkpoints"
    tensorboard_log_dir = "/mnt/DISK/xjk/pytorch_learning/tensorboard"
    log_dir = "/mnt/DISK/xjk/pytorch_learning/logs"

class LeNetConfig(Config):
    tensorboard_log_dir = Config.tensorboard_log_dir + "/" + "LeNet"
    writer = SummaryWriter(log_dir=tensorboard_log_dir)

    checkpoint_path = Config.checkpoint_dir + "/" + "checkpoint.pth.tar"
    scheduler_patience = 10
    scheduler_factor = 0.5
    valid_patience = 10

class VGGNetConfig(Config):
    train_dataset_path = "/mnt/DISK/xjk/ml/datasets/food11/training"
    valid_dataset_path = "/mnt/DISK/xjk/ml/datasets/food11/validation"
    test_dataset_path = "/mnt/DISK/xjk/ml/datasets/food11/test"

    writer_path = Config.tensorboard_log_dir + "/" + "VGGNet"
    log_path = Config.log_dir + "/" + "VGGNet"

    checkpoint_path = Config.checkpoint_dir + "/" + "VGGNet"
    scheduler_patience = 1
    scheduler_factor = 0.1
    valid_patience = 10
    momentum = 0.9
    weight_decay = 5 * 1e-4
    batch_size = 64
    kFold = 5

    full_dataset = FoodDataset.FoodDataset([
        train_dataset_path,
        valid_dataset_path
    ])
    dataset_size = len(full_dataset)




